import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import scipy.stats

trn = pd.read_csv('./all_feas_trn.csv', header=0)
tst = pd.read_csv('./all_feas_trn.csv', header=0)

def kl(x, y, z):
    k1 = scipy.stats.entropy(x, y)
    k2 = scipy.stats.entropy(x, z)
    k3 = scipy.stats.entropy(y, z)
    return k1+k2+k3


def pertile(a_in, n, drop_val=0):
    import copy
    import numpy as np
    import pandas as pd
    a = pd.Series(copy.deepcopy(a_in))
    a1 = a[a != drop_val]
    flag = np.ones(len(a1))
    b1 = copy.deepcopy(a1)
    p = np.linspace(0, 1, n + 1)
    a1p = a1.quantile(p[::-1][1:])
    for i, ta in enumerate(a1p):
        a1[(b1 >= ta) & (flag == 1)] = n - i
        flag[(b1 >= ta) & (flag == 1)] = 0
    a[a != drop_val] = a1
    return a


num = 10
atype = trn['type']
low = 5
high = 20
for i in range(trn.shape[1]):
    a = trn.iloc[:, i]
    ao = pertile(a, num)
    fea_name = ao.name
    fea_name1 = fea_name.replace('/', '__')
    if fea_name in ['type', 'ship']:
        continue

    inds = ao.unique()
    M = dict(zip(sorted(inds), range(len(inds))))
    A = pd.DataFrame({fea_name: ao, 'type': atype})
    tA = A.groupby([fea_name, 'type']).agg(len).reset_index()
    tA[fea_name] = tA[fea_name].map(M)
    df = pd.DataFrame(np.zeros((len(M), 3)), index=sorted(M.values()), columns=[0, 1, 2])
    for row in tA.iterrows():
        ind, (i, j, v) = row
        df.loc[i, j] = v
    df = df/df.sum(axis=0)
    x, y, z = df[0], df[1], df[2]
    klv = kl(x, y, z)
    print('klv: ', klv)

    plt.figure()
    plt.plot(df.index, df[0], df.index, df[1], df.index, df[2])
    plt.legend(['wei', 'tuo', 'ci'])
    plt.title(fea_name + ', kl: ' +str(klv))
    plt.savefig(f'feas/{fea_name1}.png')
    for k in [0, -1]:
        x, y, z = df.iloc[k, :]
        # if ((x <= low and y <= low and z > high) | (x <= low and z <= low and y > high) |
        #     (y <= low and z <= low and x > high)):
        if z / (y + 1) > 2 and z / (x + 1) > 2:
            print(fea_name, df)
            continue
